import cv2
import os
import numpy as np
from pathlib import Path
import shutil

class ImageMatcher:
    VALID_EXTENSIONS = {'.jpg', '.jpeg', '.png'}
    
    def __init__(self, logos_dir, target_dir, output_base_dir):
        self.invoices_dir = Path(logos_dir)
        self.target_dir = Path(target_dir)
        self.output_base_dir = Path(output_base_dir)
        self.invoices = {}
        self.load_invoices()
        
    def load_invoices(self):
        """加载所有发票模板图片"""
        for invoice_path in self.invoices_dir.glob('*'):
            if invoice_path.suffix.lower() in self.VALID_EXTENSIONS:
                invoice_img = cv2.imdecode(np.fromfile(str(invoice_path), dtype=np.uint8), cv2.IMREAD_COLOR)
                if invoice_img is not None:
                    self.invoices[invoice_path.stem] = invoice_img
        
    def match_image(self, img, invoice):
        """特征匹配函数"""
        try:
            # 转换为灰度图像可能会提高匹配效果
            img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            invoice_gray = cv2.cvtColor(invoice, cv2.COLOR_BGR2GRAY)
            
            sift = cv2.SIFT_create()
            kp1, des1 = sift.detectAndCompute(invoice_gray, None)
            kp2, des2 = sift.detectAndCompute(img_gray, None)
            
            if des1 is None or des2 is None or len(des1) < 2 or len(des2) < 2:
                return 0
            
            flann = cv2.FlannBasedMatcher(
                dict(algorithm=1, trees=5),
                dict(checks=50)
            )
            
            matches = flann.knnMatch(des1, des2, k=2)
            good_matches = [m for m, n in matches if m.distance < 0.7 * n.distance]
            return len(good_matches)
        except Exception as e:
            print(f"匹配过程出错: {str(e)}")
            return 0
        
    def process_images(self):
        """处理图片函数"""
        if not self.invoices:
            return
            
        # 创建输出目录
        for invoice_name in self.invoices:
            (self.output_base_dir / invoice_name).mkdir(parents=True, exist_ok=True)
            
        # 处理图片
        for img_path in self.target_dir.glob('*'):
            if img_path.suffix.lower() not in self.VALID_EXTENSIONS:
                continue
                
            img = cv2.imdecode(np.fromfile(str(img_path), dtype=np.uint8), cv2.IMREAD_COLOR)
            if img is None:
                continue
                
            # 寻找最佳匹配
            best_invoice = max(
                ((name, self.match_image(img, invoice))
                 for name, invoice in self.invoices.items()),
                key=lambda x: x[1],
                default=(None, 0)
            )
            
            # 如果匹配度超过阈值，复制文件
            if best_invoice[1] > 5:
                output_path = self.output_base_dir / best_invoice[0] / img_path.name
                shutil.copy2(img_path, output_path)

def main():
    dirs = {
        # 设置目录路径
        # 特征发票存放文件夹
        'invoices': r"E:\project\python\logo",
        # 待识别发票存放文件夹
        'target': r"E:\project\python\images",
        # 发票输出目录，输出后按发票类型分类。
        'output': r"E:\project\python\output"
    }
    
    try:
        # 确保输出目录存在
        os.makedirs(dirs['output'], exist_ok=True)
        
        # 检查必要目录是否存在
        if not os.path.exists(dirs['invoices']):
            print(f"发票模板目录不存在: {dirs['invoices']}")
            return
        if not os.path.exists(dirs['target']):
            print(f"待识别图片目录不存在: {dirs['target']}")
            return
            
        matcher = ImageMatcher(dirs['invoices'], dirs['target'], dirs['output'])
        matcher.process_images()
        print("处理完成")
    except Exception as e:
        print(f"程序运行出错: {str(e)}")

if __name__ == "__main__":
    main()
